import numpy as np
import matplotlib.pyplot as plt
from pyswarm import pso
from scipy.optimize import minimize, Bounds
import pandas as pd

# 常量定义
g = 9.8  # 重力加速度
missile_speed = 300  # 导弹速度
cloud_sink_speed = 3  # 云团下沉速度
cloud_effective_radius = 10  # 云团有效半径
cloud_duration = 20  # 云团有效时间

# 位置定义
fake_target = np.array([0, 0, 0])
real_target = np.array([0, 200, 0])  # 真目标位置
target_radius = 7  # 目标圆柱体半径
target_height = 10  # 目标圆柱体高度

M1_start = np.array([20000, 0, 2000])
FY1_start = np.array([17800, 0, 1800])

# 导弹方向向量（飞向假目标）
missile_dir = fake_target - M1_start
missile_dir_norm = np.linalg.norm(missile_dir)
missile_unit_dir = missile_dir / missile_dir_norm
missile_flight_time = missile_dir_norm / missile_speed  # 导弹总飞行时间


def missile_position(t):
    """计算导弹在t时刻的位置"""
    return M1_start + missile_unit_dir * missile_speed * t


def cloud_position(alpha, v, t_drop, t_burst, tau):
    """
    计算云团在起爆后tau时刻的位置
    alpha: 无人机飞行方向(弧度)
    v: 无人机速度
    t_drop: 投放时间(从任务开始)
    t_burst: 起爆时间(从投放开始)
    tau: 起爆后的时间
    """
    # 投放点位置
    drop_pos = FY1_start + np.array([v * np.cos(alpha) * t_drop,
                                     v * np.sin(alpha) * t_drop,
                                     0])

    # 起爆点位置（考虑重力影响）
    burst_pos = drop_pos + np.array([v * np.cos(alpha) * t_burst,
                                     v * np.sin(alpha) * t_burst,
                                     -0.5 * g * t_burst ** 2])

    # 云团位置（考虑下沉）
    return np.array([burst_pos[0],
                     burst_pos[1],
                     max(0, burst_pos[2] - cloud_sink_speed * tau)])  # 确保高度不小于0


def distance_to_line(point, line_start, line_end):
    """计算点到线段的距离"""
    line_vec = line_end - line_start
    point_vec = point - line_start
    line_len = np.linalg.norm(line_vec)

    if line_len < 1e-10:  # 避免除以零
        return np.linalg.norm(point - line_start)

    line_unitvec = line_vec / line_len
    point_vec_scaled = point_vec / line_len
    t = np.dot(line_unitvec, point_vec_scaled)

    if t < 0.0:
        t = 0.0
    elif t > 1.0:
        t = 1.0

    nearest = line_start + t * line_vec
    return np.linalg.norm(nearest - point)


def is_line_of_sight_blocked(missile_pos, cloud_pos, target_center, target_radius, target_height):
    """
    判断烟幕云团是否遮蔽导弹和目标圆柱体之间的视线
    烟幕中心到导弹与目标圆柱体边缘所有连线的距离的最大值小于10m即可有效遮蔽
    """
    # 生成目标圆柱体边缘的关键点
    num_points = 8
    target_edge_points = []

    # 圆柱体底部边缘
    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_radius * np.cos(angle)
        y = target_center[1] + target_radius * np.sin(angle)
        z = target_center[2]
        target_edge_points.append(np.array([x, y, z]))

    # 圆柱体顶部边缘
    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_radius * np.cos(angle)
        y = target_center[1] + target_radius * np.sin(angle)
        z = target_center[2] + target_height
        target_edge_points.append(np.array([x, y, z]))

    # 计算烟幕中心到所有导弹-目标边缘点连线的最大距离
    max_distance = 0
    for target_point in target_edge_points:
        dist = distance_to_line(cloud_pos, missile_pos, target_point)
        if dist > max_distance:
            max_distance = dist
            if max_distance > cloud_effective_radius:
                # 如果已经超过有效半径，提前返回
                return False

    # 如果最大距离小于云团有效半径，则认为视线被遮蔽
    return max_distance <= cloud_effective_radius


def precise_effective_shielding_time(alpha, v, t_drop, t_burst):
    """使用更小的时间步长计算精确的遮蔽时间（单弹）"""
    t_burst_total = t_drop + t_burst
    t_cloud_end = t_burst_total + cloud_duration

    if t_cloud_end > missile_flight_time:
        t_cloud_end = missile_flight_time

    time_step = 0.1  # 使用更小的时间步长
    effective_time = 0

    for tau in np.arange(0, min(cloud_duration, missile_flight_time - t_burst_total), time_step):
        t_current = t_burst_total + tau
        missile_pos = missile_position(t_current)
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)

        if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
            effective_time += time_step

    return effective_time


def total_effective_shielding_time(alpha, v, t_drop_list, t_burst_list, dt=0.1):
    """
    计算总有效遮蔽时间（多弹并集）
    t_drop_list: 三枚弹的投放时间列表
    t_burst_list: 三枚弹的起爆时间列表
    dt: 时间步长
    """
    T_max = missile_flight_time
    total_time = 0
    time_points = np.arange(0, T_max, dt)

    for t in time_points:
        missile_pos = missile_position(t)
        shielded = False
        for i in range(3):
            t_drop = t_drop_list[i]
            t_burst = t_burst_list[i]
            tau = t - (t_drop + t_burst)
            if tau < 0 or tau > cloud_duration:
                continue
            cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)
            if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
                shielded = True
                break
        if shielded:
            total_time += dt
    return total_time


def burst_height_constraint(x):
    """约束函数：确保起爆点高度大于0"""
    alpha, v = x[0], x[1]
    t_drop_list = [x[2], x[4], x[6]]
    t_burst_list = [x[3], x[5], x[7]]

    min_height = float('inf')
    for i in range(3):
        drop_point = FY1_start + np.array([v * np.cos(alpha) * t_drop_list[i],
                                           v * np.sin(alpha) * t_drop_list[i],
                                           0])
        burst_point = drop_point + np.array([v * np.cos(alpha) * t_burst_list[i],
                                             v * np.sin(alpha) * t_burst_list[i],
                                             -0.5 * g * t_burst_list[i] ** 2])
        min_height = min(min_height, burst_point[2])

    return min_height  # 需要大于0


def drop_interval_constraint(x):
    """约束函数：确保投放时间间隔至少1秒"""
    t_drop_list = [x[2], x[4], x[6]]
    t_drop_sorted = sorted(t_drop_list)
    min_interval = float('inf')
    for i in range(2):
        min_interval = min(min_interval, t_drop_sorted[i + 1] - t_drop_sorted[i])
    return min_interval  # 需要大于等于1


def objective_function(x):
    """
    目标函数（最大化总有效遮蔽时间）
    x = [alpha, v, t_drop1, t_burst1, t_drop2, t_burst2, t_drop3, t_burst3]
    """
    alpha, v = x[0], x[1]
    t_drop_list = [x[2], x[4], x[6]]
    t_burst_list = [x[3], x[5], x[7]]

    # 约束检查
    if v < 70 or v > 140:
        return 1e9
    if alpha < 0 or alpha > 2 * np.pi:
        return 1e9
    for i in range(3):
        if t_drop_list[i] < 0 or t_burst_list[i] < 0:
            return 1e9
        if t_drop_list[i] + t_burst_list[i] > missile_flight_time:
            return 1e9

    # 检查起爆点高度
    for i in range(3):
        drop_point = FY1_start + np.array([v * np.cos(alpha) * t_drop_list[i],
                                           v * np.sin(alpha) * t_drop_list[i],
                                           0])
        burst_point = drop_point + np.array([v * np.cos(alpha) * t_burst_list[i],
                                             v * np.sin(alpha) * t_burst_list[i],
                                             -0.5 * g * t_burst_list[i] ** 2])
        if burst_point[2] < 0:
            return 1e9

    # 检查投放时间间隔
    t_drop_sorted = sorted(t_drop_list)
    for i in range(2):
        if t_drop_sorted[i + 1] - t_drop_sorted[i] < 1.0:
            return 1e9

    shielding_time = total_effective_shielding_time(alpha, v, t_drop_list, t_burst_list, dt=0.1)

    # 添加惩罚项：鼓励三枚弹都起作用
    individual_times = [precise_effective_shielding_time(alpha, v, t_drop_list[i], t_burst_list[i]) for i in range(3)]
    zero_count = sum(1 for t in individual_times if t < 0.1)
    penalty = zero_count * 20  # 每有一枚弹不起作用，惩罚20秒

    return -(shielding_time - penalty)


# PSO参数范围 - 基于第二问的结果进行调整
lb = [0, 70, 0, 0, 10, 0, 20, 0]  # 下界
ub = [2 * np.pi, 140, 30, 20, 40, 20, 60, 20]  # 上界

# 设置随机种子以确保结果可重复
np.random.seed(42)

print("开始PSO优化...")
# 运行PSO优化
xopt, fopt = pso(objective_function, lb, ub, swarmsize=100, maxiter=200, minstep=1e-4)

print("PSO优化完成，开始局部搜索...")
# 使用局部搜索进一步优化
bounds = Bounds(lb, ub)
constraints = [
    {'type': 'ineq', 'fun': burst_height_constraint},
    {'type': 'ineq', 'fun': drop_interval_constraint}
]
result = minimize(objective_function, xopt, method='SLSQP', bounds=bounds, constraints=constraints)
xopt_refined = result.x
fopt_refined = result.fun

# 提取优化参数
alpha_opt, v_opt = xopt_refined[0], xopt_refined[1]
t_drop_opt = [xopt_refined[2], xopt_refined[4], xopt_refined[6]]
t_burst_opt = [xopt_refined[3], xopt_refined[5], xopt_refined[7]]

# 计算精确的总遮蔽时间（并集）
precise_total_shielding_time = total_effective_shielding_time(alpha_opt, v_opt, t_drop_opt, t_burst_opt, dt=0.1)

# 计算每枚弹的单独遮蔽时间
shielding_times = []
for i in range(3):
    time_i = precise_effective_shielding_time(alpha_opt, v_opt, t_drop_opt[i], t_burst_opt[i])
    shielding_times.append(time_i)

# 计算投放点和起爆点坐标
drop_points = []
burst_points = []
for i in range(3):
    drop_point = FY1_start + np.array([v_opt * np.cos(alpha_opt) * t_drop_opt[i],
                                       v_opt * np.sin(alpha_opt) * t_drop_opt[i],
                                       0])
    burst_point = drop_point + np.array([v_opt * np.cos(alpha_opt) * t_burst_opt[i],
                                         v_opt * np.sin(alpha_opt) * t_burst_opt[i],
                                         -0.5 * g * t_burst_opt[i] ** 2])
    drop_points.append(drop_point)
    burst_points.append(burst_point)

# 输出结果
print("优化结果:")
print(f"飞行方向: {np.degrees(alpha_opt):.2f}°")
print(f"飞行速度: {v_opt:.2f} m/s")
print(f"总有效遮蔽时间: {precise_total_shielding_time:.2f} s")
for i in range(3):
    print(
        f"弹 {i + 1}: 投放时间={t_drop_opt[i]:.2f}s, 起爆时间={t_burst_opt[i]:.2f}s, 单独遮蔽时间={shielding_times[i]:.2f}s")
    print(f"投放点坐标: ({drop_points[i][0]:.2f}, {drop_points[i][1]:.2f}, {drop_points[i][2]:.2f})")
    print(f"起爆点坐标: ({burst_points[i][0]:.2f}, {burst_points[i][1]:.2f}, {burst_points[i][2]:.2f})")

# 保存到result.xlsx
data = {
    '无人机运动方向': [np.degrees(alpha_opt)] * 3,
    '无人机运动速度 (m/s)': [v_opt] * 3,
    '烟幕干扰弹编号': [1, 2, 3],
    '烟幕干扰弹投放点的x坐标 (m)': [drop_points[i][0] for i in range(3)],
    '烟幕干扰弹投放点的y坐标 (m)': [drop_points[i][1] for i in range(3)],
    '烟幕干扰弹投放点的z坐标 (m)': [drop_points[i][2] for i in range(3)],
    '烟幕干扰弹起爆点的x坐标 (m)': [burst_points[i][0] for i in range(3)],
    '烟幕干扰弹起爆点的y坐标 (m)': [burst_points[i][1] for i in range(3)],
    '烟幕干扰弹起爆点的z坐标 (m)': [burst_points[i][2] for i in range(3)],
    '有效干扰时长 (s)': shielding_times
}

df = pd.DataFrame(data)
df.to_excel('result.xlsx', index=False)
print("结果已保存到result.xlsx")


# 可视化遮蔽效果
def plot_shielding_effect(alpha, v, t_drop_list, t_burst_list):
    """可视化遮蔽效果"""
    # 字体和线条设置
    plt.rcParams['font.sans-serif'] = ['PingFang SC', 'Songti SC', 'Arial Unicode MS']
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

    # 调整字体大小
    plt.rcParams['font.size'] = 14  # 全局字体大小
    plt.rcParams['axes.titlesize'] = 16  # 标题字体大小
    plt.rcParams['axes.labelsize'] = 15  # 轴标签字体大小
    plt.rcParams['xtick.labelsize'] = 13  # x轴刻度字体大小
    plt.rcParams['ytick.labelsize'] = 13  # y轴刻度字体大小
    plt.rcParams['legend.fontsize'] = 14  # 图例字体大小

    T_max = missile_flight_time
    dt = 0.1
    time_points = np.arange(0, T_max, dt)
    shielding_status = np.zeros_like(time_points)

    for idx, t in enumerate(time_points):
        missile_pos = missile_position(t)
        shielded = False
        for i in range(3):
            t_drop = t_drop_list[i]
            t_burst = t_burst_list[i]
            tau = t - (t_drop + t_burst)
            if tau < 0 or tau > cloud_duration:
                continue
            cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)
            if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
                shielded = True
                break
        shielding_status[idx] = 1 if shielded else 0

    plt.figure(figsize=(12, 6))
    plt.plot(time_points, shielding_status, 'b-', linewidth=2.0, label='总遮蔽状态')

    # 绘制每枚弹的单独遮蔽状态
    for i in range(3):
        individual_status = np.zeros_like(time_points)
        for idx, t in enumerate(time_points):
            t_drop = t_drop_list[i]
            t_burst = t_burst_list[i]
            tau = t - (t_drop + t_burst)
            if tau < 0 or tau > cloud_duration:
                continue
            missile_pos = missile_position(t)
            cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)
            if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
                individual_status[idx] = 1

        plt.plot(time_points, individual_status, '--', linewidth=1.5, label=f'弹 {i + 1} 遮蔽状态')

    plt.xlabel('时间 (s)')
    plt.ylabel('遮蔽状态 (1=遮蔽, 0=未遮蔽)')
    plt.title('导弹与目标之间视线遮蔽状态随时间变化')
    plt.legend()
    plt.grid(True, linewidth=1.0)
    plt.show()


print("\n生成遮蔽效果图...")
plot_shielding_effect(alpha_opt, v_opt, t_drop_opt, t_burst_opt)